package middleware

import (
	"context"
	"net/http"
	"strings"

	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest"
	"github.com/zeromicro/go-zero/rest/httpx"
	"searchcenter/internal/utils" // 替换为你的实际路径
)

type contextKey string

const (
	UserIDKey contextKey = "userId" // 常量命名建议使用驼峰（如UserIDKey）
	TokenKey  contextKey = "token" // 新增
)

// NewAuthMiddleware 创建JWT认证中间件
// secret: JWT加密密钥
func NewAuthMiddleware(secret string) rest.Middleware {
	return func(next http.HandlerFunc) http.HandlerFunc {
		return func(w http.ResponseWriter, r *http.Request) {
			// 1. 获取Token
			authHeader := r.Header.Get("Authorization")
			if authHeader == "" {
				httpx.WriteJson(w, http.StatusUnauthorized, map[string]interface{}{
					"code":    401,
					"message": "未提供认证Token",
				})
				return
			}

			// 2. 提取Bearer Token（更健壮的检查）
			const bearerPrefix = "Bearer "
			if !strings.HasPrefix(authHeader, bearerPrefix) {
				httpx.WriteJson(w, http.StatusUnauthorized, map[string]interface{}{
					"code":    401,
					"message": "Token格式错误，需要Bearer Token",
				})
				return
			}
			token := strings.TrimPrefix(authHeader, bearerPrefix)

			// 3. 解析Token
			claims, err := utils.ParseToken(token, secret)
			if err != nil {
				logx.Errorf("解析Token失败: %v", err)
				httpx.WriteJson(w, http.StatusUnauthorized, map[string]interface{}{
					"code":    401,
					"message": "无效的Token: " + err.Error(),
				})
				return
			}

			// 4. 存储用户ID到Context（避免Key冲突，使用自定义类型）
			ctx := context.WithValue(r.Context(), UserIDKey, claims.UserId)
			ctx = context.WithValue(ctx, TokenKey, token)
			next(w, r.WithContext(ctx))
		}
	}
}